##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)
library(ggpubr)

##########################################################################################

option_list <- list(
    make_option(c("--info_file"), type = "character") ,
    make_option(c("--maf_file"), type = "character") ,
    make_option(c("--maf_msi_file"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    info_file <- paste(work_dir,"/config/STAD-useCombine.Sample.tsv",sep="")
    maf_file <- paste(work_dir,"/maf/All_ForMutBurden.extract.maf",sep="")
    maf_msi_file <- paste(work_dir,"/maf/All_ForMutBurden.extract.MSI.maf",sep="")
	images_path <- paste(work_dir,"/images/mutBurden",sep="")

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

maf_file <- opt$maf_file
maf_msi_file <- opt$maf_msi_file
info_file <- opt$info_file
images_path <- opt$images_path

dir.create(images_path , recursive = T)

###########################################################################################

if(1!=1){
    col <- c(
      brewer.pal(9,"YlGnBu")[6],
      rgb(234,106,79,alpha=255,maxColorValue=255),
      rgb(203,24,30,alpha=255,maxColorValue=255),
      rgb(255,0,0,alpha=255,maxColorValue=255)
      )
}
col <- c(
    brewer.pal(9,"YlGnBu")[6],
    rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
    rgb(red=2,green=100,blue=190,alpha=255,max=255) ,
    rgb(255,0,0,alpha=255,maxColorValue=255)
    )


names(col) <- c("IM" , "IGC" , "DGC" , "GC")

col_im <- brewer.pal(9,"YlGnBu")[6:8]
names(col_im) <- c("IM(IGC)" , "IM(DGC)" , "IM(IGC_DGC)")

Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

###########################################################################################

dat_sample <- data.frame(fread( info_file ))
dat_maf <- data.frame(fread( maf_file ))
dat_maf_msi <- data.frame(fread( maf_msi_file ))

###########################################################################################
dat_maf$MS_Type <- "MSS"
dat_maf_msi$MS_Type <- "MSI"

dat_maf <- rbind(dat_maf , dat_maf_msi)

###########################################################################################
## 计算WES的突变数量
dat_maf <- subset( dat_maf , Variant_Classification %in% Variant_Types )
exon_mut <- dat_maf %>%
group_by( Tumor_Sample_Barcode , MS_Type ) %>%
summarize( exonMutNum = length(Start_position) )

dat_plot <- merge( dat_sample , exon_mut , by.x = "Tumor" , by.y = "Tumor_Sample_Barcode" , all.x = TRUE )
dat_plot$BurdenAll <- dat_plot$mutNum/(dat_plot$coverage_All/1024/1024)
dat_plot$BurdenExon <- dat_plot$exonMutNum/(dat_plot$coverage_CDS/1024/1024)
dat_plot$BurdenAll[is.na(dat_plot$BurdenAll)] <- 0
dat_plot$BurdenExon[is.na(dat_plot$BurdenExon)] <- 0

## 部分肠化样本无外显子突变
for(id in dat_plot[is.na(dat_plot$MS_Type),"Patient"]){
    ms_type <- dat_plot[dat_plot$Patient==id,"MS_Type"][!is.na(dat_plot[dat_plot$Patient==id,"MS_Type"])]
    dat_plot[dat_plot$Patient==id,"MS_Type"] <- unique(ms_type)
}

out_name <- paste0( images_path , "/mutBurden.tsv" )
write.table( dat_plot , out_name , row.names = F , sep = '\t' , quote = F  )

###########################################################################################
## 多个样本负荷取中位数
dat_plot2 <- dat_plot %>%
group_by( Patient , Class , Type , MS_Type ) %>%
summarize( BurdenAll = median(BurdenAll) ,BurdenExon = median(BurdenExon) )

## IGC和DGC合并
dat_plot_gc <- subset( dat_plot , Class %in% c("IGC" , "DGC") )
dat_plot_gc$Class <- "GC"
dat_plot_gc$Type <- "IM + GC"
dat_plot3 <- dat_plot_gc %>%
group_by( Patient , Class , Type , MS_Type  ) %>%
summarize( BurdenAll = median(BurdenAll) ,BurdenExon = median(BurdenExon) )

dat_plot4 <- dat_plot2
dat_plot4$Type <- "All"

dat_plot_final <- rbind( dat_plot2 , dat_plot3 , dat_plot4 )
dat_plot_final$Class <- factor( dat_plot_final$Class , levels = names(col) , order = T )

dat_plot4_out <- dat_plot4 %>%
group_by( Class , MS_Type ) %>%
summarize( BurdenAll = median(BurdenAll) , BurdenExon = median(BurdenExon) )

out_name <- paste0( images_path , "/mutBurden.class.tsv" )
write.table( dat_plot4_out , out_name , row.names = F , sep = '\t' , quote = F  )

###########################################################################################

plotBurden <- function(dat_plot_tmp = dat_plot_tmp , mutType = mutType){

    my_comparisons_1 <- list( 
    c(1, 2) , c(1, 3) ,
    c(2, 3)
    )

    plot <- ggplot( dat_plot_tmp , aes( x = Class , y = MutBurden_use , color = Class ) ) +
        geom_line( aes( group = Patient ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
        geom_boxplot(alpha =1 , size = 0.9 , width = 0.6 , outlier.shape = NA) +
        geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
        scale_color_manual(values=col) +
        xlab(NULL) +
        #scale_y_continuous(trans="sqrt" , breaks = c(0,0.5,1,2,3,4,5,10,20,30,40,50) ) +
        ylab("Mutation rate per MB")+
        theme_bw() +
        stat_compare_means(
            comparisons = my_comparisons_1 , 
            method = "wilcox.test" , 
            aes(as_label = ifelse(p < 0.05,sprintf("p = %2.1e", as.numeric(..p.format..)), ..p.format..))
            ) +
        theme(panel.background = element_blank(),#设置背影为白色#清除网格线
            legend.position ='left',
            legend.title = element_blank() ,
            panel.grid.major=element_line(colour=NA),
            plot.title = element_text(size = 8,color="black",face='bold'),
            legend.text = element_text(size = 10,color="black",face='bold'),
            axis.text.y = element_text(size = 10,color="black",face='bold'),
            axis.title.x = element_text(size = 10,color="black",face='bold'),
            axis.title.y = element_text(size = 10,color="black",face='bold'),
            axis.ticks.x = element_blank(),
            axis.text.x = element_text(size = 10,color="black",face='bold') ,
            axis.line = element_line(size = 0.5)) 
    
    out_name <- paste0( images_path , "/mutBurden.oneImage." , mutType , ".pdf" )  
    ggsave(file=out_name,plot=plot,width=5,height=5)

}

## 排除既存在IGC和DGC的患者
dat_plot_tmp <- subset( dat_plot2 , Type != "IM + IGC + DGC" )
dat_plot_tmp$Class <- factor( dat_plot_tmp$Class , levels = names(col) , order = T )
dat_plot_tmp$MutBurden_use <- dat_plot_tmp$BurdenExon

## 分MSI和MSS分布比较
for( msi in unique(unique(dat_plot_tmp$MS_Type)) ){
    dat_plot_tmp_use <- subset( dat_plot_tmp , MS_Type == msi )
    mutType <- msi
    plotBurden(dat_plot_tmp = dat_plot_tmp_use , mutType = mutType)
}


###########################################################################################

trans <- function(num){
    up <- floor(log10(num))
    down <- round(num*10^(-up),2)
    text <- paste("P == ",down," %*% 10","^",up)
    return(text)
}

###########################################################################################
## 分IM+IGC、IM+DGC分别比较
dat_plot_tmp <- data.frame(dat_plot2)
dat_plot_tmp <- subset( dat_plot_tmp , Type!="IM + IGC + DGC")
dat_plot_tmp$Class_use <- ifelse(dat_plot_tmp$Class == "IM" , "IM" , "GC")
dat_plot_tmp$Class_use <- factor( dat_plot_tmp$Class_use , levels = c("IM" , "GC") , order = T )
dat_plot_tmp$Type <- factor( dat_plot_tmp$Type , levels = c("IM + IGC" , "IM + DGC") , order = T )

my_comparisons_1 <- list( c(1, 2) )

for( msi in unique(unique(dat_plot_tmp$MS_Type)) ){
    
    dat_plot_tmp_use <- subset( dat_plot_tmp , MS_Type == msi )
    type_num <- dat_plot_tmp_use %>%
    group_by( Type ) %>%
    summarize( type_nums = length(unique(Patient)) )
    dat_plot_tmp_use <- merge( dat_plot_tmp_use , type_num , by = "Type" )
    #dat_plot_tmp_use$Type <- paste0( dat_plot_tmp_use$Type , "\n" , "(" , dat_plot_tmp_use$type_nums , ")" )

    dat_tmp <- c()

    y_max <- max(dat_plot_tmp_use$BurdenExon) + 3

    for( type in unique(dat_plot_tmp_use$Type) ){

        dat <- subset( dat_plot_tmp_use , Type == type )
        dat <- dat[order(dat$Patient),]

        a <- dat[dat$Class==unique(dat$Class)[1],"BurdenExon"]
        b <- dat[dat$Class==unique(dat$Class)[2],"BurdenExon"]
        
        p <- wilcox.test( a , b , paired = T )$p.value

        if( p < 0.01 ){
            p_text <- trans(p)
        }else{
            p_text <- paste0( "P == " , round(as.numeric(p) , 3) ) 
        }
        dat$p_text <- ""
        dat$p_text[1] <- p_text
        dat_tmp <- rbind(dat_tmp , dat)
    }
    dat_tmp$Class <- factor( dat_tmp$Class , levels = c("IM" , "IGC" , "DGC") , order = T )
    #dat_tmp$Type <- factor( dat_tmp$Type , levels = c("IM + IGC\n(50)" , "IM + DGC\n(31)") , order = T )
    dat_tmp$Type <- factor( dat_tmp$Type , levels = c("IM + IGC" , "IM + DGC") , order = T )

    col_tmp <- c(
        rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
        rgb(red=2,green=100,blue=190,alpha=255,max=255) ,
        rgb(red=2,green=100,blue=190,alpha=255,max=255) 
    )

    names(col_tmp) <- c("IM" , "IGC" , "DGC")

    plot <- ggplot( dat_tmp , aes( x = Class , y = BurdenExon , color = Class , fill = Class ) ) +
        geom_line( aes( group = Patient ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
        #geom_boxplot(size = 1.2 , outlier.shape = NA ) + ## 去除散点，加粗线
        #geom_jitter(position = position_jitterdodge(0.8) , size = 1) + 
        geom_violin(trim=FALSE) +
        geom_boxplot(width=0.2,position=position_dodge(0.9),fill="white",color="black")+ #绘制箱线图
        scale_color_manual(values=col_tmp) +
        scale_fill_manual(values=col_tmp) +
        facet_grid(.~Type,space='free_x',scales='free_x') +
        xlab(NULL) +
        #scale_y_sqrt() +
        ylab("Mutation rate per MB")+
        theme_bw() +
        geom_text(aes(label=p_text , y = y_max , x = 1.5),parse = TRUE,size=5 , color = "black" , face='bold') +
        #stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test") +
        theme(
            legend.position = 'none',
            legend.title = element_blank() ,
            panel.grid.major=element_blank(),
            panel.grid.minor=element_blank(),
            panel.background = element_blank(),
            panel.border = element_blank(),
            plot.title = element_text(size = 12,color="black",face='bold'),
            legend.text = element_text(size = 12,color="black",face='bold'),
            axis.text.y = element_text(size = 15,color="black",face='bold'),
            axis.title.x = element_text(size = 15,color="black",face='bold'),
            axis.title.y = element_text(size = 12,color="black",face='bold'),
            axis.text.x = element_text(size = 15,color="black",face='bold') ,
            axis.ticks.length = unit(0.2, "cm") ,
            strip.text.x = element_text(size = 15, colour = "black",face='bold') ,
            axis.line = element_line(size = 0.5)) 
    out_name <- paste0( images_path , "/mutBurden.Type.IM_IGC-IM_DGC.",msi,".pdf" )  
    ggsave(file=out_name,plot=plot,width=4.7/1.2,height=4.8/1.2)

    tmp_res <- dat_tmp %>%
    group_by( Type , Class ) %>%
    summarize( BurdenExon = median(BurdenExon) )
    out_name <- paste0( images_path , "/mutBurden.Type.IM_IGC-IM_DGC.",msi,".tsv" )  
    write.table( tmp_res , out_name , row.names = F , sep = "\t" , quote = F )

}

###########################################################################################
## IM + IGC + DGC的
dat_plot_tmp <- data.frame(dat_plot2)
dat_plot_tmp <- subset( dat_plot_tmp , Type=="IM + IGC + DGC")
dat_plot_tmp$Class <- factor( dat_plot_tmp$Class , levels = c("IM" , "IGC" , "DGC") , order = T )

my_comparisons_1 <- list( c(1, 2) , c(1,3) , c(2,3) )

plot <- ggplot( dat_plot_tmp , aes( x = Class , y = BurdenExon , color = Class ) ) +
    geom_line( aes( group = Patient ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
    geom_boxplot(alpha =1 , size = 0.9 , width = 0.6 , outlier.shape = NA) +
    geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
    scale_color_manual(values=col) +
    facet_grid(.~Type,space='free_x',scales='free_x') +
    ylim(0,15)+
    xlab(NULL) +
    ylab("Mutation rate per MB")+
    theme_bw() +
    stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test") +
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='left',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        plot.title = element_text(size = 8,color="black",face='bold'),
        legend.text = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 10,color="black",face='bold'),
        axis.title.x = element_text(size = 10,color="black",face='bold'),
        axis.title.y = element_text(size = 10,color="black",face='bold'),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(size = 10,color="black",face='bold') ,
        axis.line = element_line(size = 0.5)) 

out_name <- paste0( images_path , "/mutBurden.Type.IM-IGC-DGC.pdf" )  
ggsave(file=out_name,plot=plot,width=4,height=5)

###########################################################################################
## 不同病理亚型的IM
dat_plot_tmp <- data.frame(dat_plot2)
dat_plot_tmp <- subset( dat_plot_tmp , Class == "IM" )
dat_plot_tmp <- subset( dat_plot_tmp , Type != "IM + IGC + DGC")
dat_plot_tmp$Class <- ifelse( dat_plot_tmp$Type == "IM + IGC" , "IM(IGC)" , dat_plot_tmp$Class )
dat_plot_tmp$Class <- ifelse( dat_plot_tmp$Type == "IM + DGC" , "IM(DGC)" , dat_plot_tmp$Class )
#dat_plot_tmp$Class <- ifelse( dat_plot_tmp$Type == "IM + IGC + DGC" , "IM(IGC_DGC)" , dat_plot_tmp$Class )
#dat_plot_tmp$Class <- factor( dat_plot_tmp$Class , levels = c("IM(IGC)" , "IM(DGC)" , "IM(IGC_DGC)") , order = T )
dat_plot_tmp$Class <- factor( dat_plot_tmp$Class , levels = c("IM(IGC)" , "IM(DGC)") , order = T )

for( msi in unique(unique(dat_plot_tmp$MS_Type)) ){

    dat_plot_tmp_use <- subset( dat_plot_tmp , MS_Type == msi )
    dat_tmp <- c()

    y_max <- max(dat_plot_tmp_use$BurdenExon) + 1

    dat <- dat_plot_tmp_use
    a <- dat[dat$Class==unique(dat$Class)[1],"BurdenExon"]
    b <- dat[dat$Class==unique(dat$Class)[2],"BurdenExon"]
    p <- wilcox.test( a , b )$p.value

    if( p < 0.01 ){
        p_text <- trans(p)
    }else{
        p_text <- paste0( "p == " , round(as.numeric(p) , 3) ) 
    }
    dat$p_text <- ""
    dat$p_text[1] <- p_text
    dat_tmp <- rbind(dat_tmp , dat)

    col_tmp <- c(
        rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
        rgb(red=2,green=100,blue=190,alpha=255,max=255) 
    )

    names(col_tmp) <- c("IM(IGC)" , "IM(DGC)")

    plot <- ggplot( dat_tmp , aes( x = Class , y = BurdenExon , color = Class ) ) +
        geom_boxplot(alpha =1 , size = 0.9 , width = 0.6 , outlier.shape = NA) +
        geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
        scale_color_manual(values=col_tmp) +
        xlab(NULL) +
        ylab("Mutation rate per MB")+
        theme_bw() +
        geom_text(aes(label=p_text , y = y_max , x = 1.5),parse = TRUE,size=4 , color = "black") +
        theme(panel.background = element_blank(),#设置背影为白色#清除网格线
            legend.position ='none',
            legend.title = element_blank() ,
            panel.grid.major=element_line(colour=NA),
            plot.title = element_text(size = 8,color="black",face='bold'),
            legend.text = element_text(size = 10,color="black",face='bold'),
            axis.text.y = element_text(size = 10,color="black",face='bold'),
            axis.title.x = element_text(size = 10,color="black",face='bold'),
            axis.title.y = element_text(size = 10,color="black",face='bold'),
            axis.text.x = element_text(size = 10,color="black",face='bold') ,
            axis.line = element_line(size = 0.5)) 

    out_name <- paste0( images_path , "/mutBurden.Type.IM.",msi,".pdf" )  
    ggsave(file=out_name,plot=plot,width=2,height=3)
}

###########################################################################################
## MSI的IM和MSS的IM
dat_plot_tmp <- data.frame(dat_plot2)
dat_plot_tmp <- subset( dat_plot_tmp , Class == "IM" )

my_comparisons_1 <- list( c(1, 2) )

plot <- ggplot( dat_plot_tmp , aes( x = MS_Type , y = BurdenExon , color = MS_Type ) ) +
    geom_boxplot(alpha =1 , size = 0.9 , width = 0.6 , outlier.shape = NA) +
    geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
    scale_color_manual(values=as.character(col_im)) +
    xlab(NULL) +
    ylab("Mutation rate per MB")+
    theme_bw() +
    stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test") +
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='left',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        plot.title = element_text(size = 8,color="black",face='bold'),
        legend.text = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 10,color="black",face='bold'),
        axis.title.x = element_text(size = 10,color="black",face='bold'),
        axis.title.y = element_text(size = 10,color="black",face='bold'),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(size = 10,color="black",face='bold') ,
        axis.line = element_line(size = 0.5)) 

out_name <- paste0( images_path , "/mutBurden.Type.IM.MSI_MSS.pdf" )  
ggsave(file=out_name,plot=plot,width=4,height=5)

###########################################################################################
## MSI的IM和MSS的IM
## 不同亚型
dat_plot_tmp <- data.frame(dat_plot2)
dat_plot_tmp <- subset( dat_plot_tmp , Class == "IM" )

my_comparisons_1 <- list( c(1, 2) )

plot <- ggplot( dat_plot_tmp , aes( x = MS_Type , y = BurdenExon , color = MS_Type ) ) +
    geom_boxplot(alpha =1 , size = 0.9 , width = 0.6 , outlier.shape = NA) +
    geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
    scale_color_manual(values=as.character(col_im)) +
    facet_grid(.~Type,space='free_x',scales='free_x') +
    xlab(NULL) +
    ylab("Mutation rate per MB")+
    theme_bw() +
    stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test") +
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='left',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        plot.title = element_text(size = 8,color="black",face='bold'),
        legend.text = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 10,color="black",face='bold'),
        axis.title.x = element_text(size = 10,color="black",face='bold'),
        axis.title.y = element_text(size = 10,color="black",face='bold'),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(size = 10,color="black",face='bold') ,
        axis.line = element_line(size = 0.5)) 

out_name <- paste0( images_path , "/mutBurden.Type.IM.MSI_MSS.divide.pdf" )  
ggsave(file=out_name,plot=plot,width=5,height=5)